#!/bin/bash

set -e
set -o pipefail

echo The content of the submitted script:
echo ===========================================
cat "${0}"
echo ===========================================

echo
echo Runing environment:
echo ===========================================
echo "# work_dir ----------------->  $(pwd)   "
echo "# command line arguments --->           "
count=0
for var in "$@"; do
    ((count += 1))
    echo "# \$$count ----------------->  $var   "
done
echo "# running script path ------>  ${0}     "
echo "# script location ---------->  ${0%/*}  "
echo "# script name -------------->  ${0##*/} "
echo ===========================================
echo

echo Slurm environment variables:
echo ===========================================
echo "SLURM_JOBID = $SLURM_JOBID"
echo "SLURM_JOB_NAME = $SLURM_JOB_NAME"
echo "SLURM_JOB_PARTITION = $SLURM_JOB_PARTITION"

echo "SLURM_JOB_NODELIST = $SLURM_JOB_NODELIST"
echo "SLURM_NNODES = $SLURM_NNODES"
echo "SLURM_NTASKS=$SLURM_NTASKS"
echo "SLURM_NTASKS_PER_CORE=$SLURM_NTASKS_PER_CORE"
echo "SLURM_NTASKS_PER_NODE=$SLURM_NTASKS_PER_NODE"
echo "SLURM_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK"
echo "SLURM_SUBMIT_DIR=$SLURM_SUBMIT_DIR"
echo ===========================================
echo

for i in /etc/profile.d/*.sh /etc/profile.d/sh.local; do
    if [ -r "$i" ]; then
        if [ "${-#*i}" != "$-" ]; then
            # shellcheck disable=SC1090
            . "$i"
        else
            # shellcheck disable=SC1090
            . "$i" >/dev/null
        fi
    fi
done

for key in "$@"; do
    case $key in
    --input_dir=*)
        input_dir="${key#*=}"
        shift # past argument=value
        ;;
    --output_dir=*)
        output_dir="${key#*=}"
        shift # past argument=value
        ;;
    --work_dir=*)
        work_dir="${key#*=}"
        shift # past argument=value
        ;;
    --auxiliary_dir=*)
        auxiliary_dir="${key#*=}"
        shift # past argument=value
        ;;
    *) ;;

    esac
done

echo "------------------------------------------"
echo Task Start

if [ -z "$input_dir" ]; then
    echo "please input valid input file"
    exit 1
fi

if [ -z "$output_dir" ]; then
    echo "please input valid output directory"
    exit 1
fi

if [ -z "$work_dir" ]; then
    echo "please input valid work directory"
    exit 1
fi

if [ -n "$auxiliary_dir" ] && [ ! -d "$auxiliary_dir" ]; then
    echo "Auxiliary directory $auxiliary_dir not exists"
    exit 1
fi

echo "input_dir=""$input_dir"
echo "output_dir=""$output_dir"
echo "work_dir=$work_dir"
echo "auxiliary_dir=$auxiliary_dir"

mkdir -p "$output_dir"
# Make a subdirectory and run GTC program in it.
# Because GTC outputs to its working directory,
# it will be easier to copy the output to the destination
# output_dir if we are working in a clean subdirectory.
mkdir -p "$work_dir/$SLURM_JOB_NAME"
cd "$work_dir/$SLURM_JOB_NAME"

## Modules: ====> Loading desired module (DO NOT MODIFY)
module load gtc/3.21

# 解决报错：
#  MPI startup(): PMI server not found. Please set I_MPI_PMI_LIBRARY variable if it is not a singleton case.
# see: https://bugs.schedmd.com/show_bug.cgi?id=6727
echo I_MPI_PMI_LIBRARY = "$I_MPI_PMI_LIBRARY"
if [ -z "$I_MPI_PMI_LIBRARY" ]; then
    export I_MPI_PMI_LIBRARY=/usr/lib64/libpmi.so
fi

echo -e "[Star time]\t$(date +"%Y-%m-%d %H:%M:%S")\t"

cp -r "$input_dir" ./conf

if [ -d "$auxiliary_dir" ]; then
    if [ -f "$auxiliary_dir/auxiliary_job.cpp" ]; then
        echo "Find source file of the auxiliary job. "
        cp "$auxiliary_dir/auxiliary_job.cpp" .
        mpiicpc auxiliary_job.cpp -qopenmp -o auxiliary_job
        chmod +x auxiliary_job
    else
        echo "No auxiliary program found. Ignore it."
    fi

    if [ -f "$auxiliary_dir/start-gtc.sh" ]; then
        echo "Find user specified start script. "
        cp "$auxiliary_dir"/start-gtc.sh .
        chmod +x start-gtc.sh
    else
        echo "No start script found. Use default instead."
    fi
fi

if [ ! -f start-gtc.sh ]; then
    echo '#!/bin/bash' >start-gtc.sh
    if [ -x ./auxiliary_job ]; then
        echo './auxiliary_job' >>start-gtc.sh
    fi
    echo 'gtc' >>start-gtc.sh
fi

echo contents of start script:
echo ===========================================
cat start-gtc.sh
echo ===========================================

chmod +x start-gtc.sh

ulimit -s unlimited
mkdir -p restart_dir{1,2} {phi,trackp}_dir

echo
echo Final Environment:
echo ===========================================
echo "# user --------------------->  $(whoami)"
echo "# submit host -------------->  $(hostname)"
echo "# work_dir ----------------->  $(pwd)"
echo "# input_dir ---------------->  ${input_dir}  "
echo "# output_dir --------------->  ${output_dir} "
echo "# gtc executable location -->  $(which gtc)"
echo "# contents of work_dir ----->  "
ls -R -l
echo "# 'ulimit -a' output ------->  "
ulimit -a
if [ -f /.dockerenv ]; then
    echo "# world -------------------->  docker"
elif [ -f /proc/1/cgroup ] && grep -sq 'docker\|lxc' /proc/1/cgroup; then
    echo "# world -------------------->  docker/lxc"
else
    echo "# world -------------------->  real"
fi
echo ===========================================
echo

srun --cpus-per-task 1 ./start-gtc.sh

mv ./* "$output_dir"/

echo -e "[End time]\t$(date +"%Y-%m-%d %H:%M:%S")\t"
echo "Run Finished"
